#include"bmp.h"
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>


// #include<iostream>
 
// using namespace std;

void YUV420P2RGB32(const uchar *yuv_buffer_in, const uchar *rgb_buffer_out, int width, int height);


int AVFrame2Img(AVFrame *frame, uchar *rgb_data) {
    // 获取帧宽、高及通道数量
    int frame_height = frame->height;
    int frame_width = frame->width;
    int frame_channels = 3;
	int image_size  = frame_height * frame_width * sizeof(uchar)*frame_channels;

    // 初始化 Mat
    // Mat img = Mat::zeros(frame_height, frame_width, CV_8UC3);
	// uchar* rgb_data = (uchar*)malloc(image_size);
    // 初始化存放 YUV 编码图片的 buffer 内存空间
    uchar* yuv_buffer = (uchar*)malloc(image_size);

    // 获取图片原始数据
    // Y
    for (int i = 0; i < frame_height; i++) {
        memcpy(yuv_buffer + frame_width * i,
               frame->data[0] + frame->linesize[0] * i,
               frame_width);
    }
    // U
    for (int j = 0; j < frame_height / 2; j++) {
        memcpy(yuv_buffer + frame_width * frame_height + frame_width / 2 * j,
               frame->data[1] + frame->linesize[1] * j,
               frame_width / 2);
    }
    // V
    for (int k = 0; k < frame_height / 2; k++) {
        memcpy(yuv_buffer + frame_width * frame_height + frame_width / 2 * (frame_height / 2) + frame_width / 2 * k,
               frame->data[2] + frame->linesize[2] * k,
               frame_width / 2);
    }

    // 转换为 RGB 编码
    YUV420P2RGB32(yuv_buffer, rgb_data, frame_width, frame_height);

    // 释放 buffer 内存空间
    free(yuv_buffer);
    return 0;
}

void YUV420P2RGB32(const uchar *yuv_buffer_in, const uchar *rgb_buffer_out, int width, int height) {
    uchar *yuv_buffer = (uchar *)yuv_buffer_in;
    uchar *rgb_buffer = (uchar *)rgb_buffer_out;

    int channels = 3;

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            int index_Y = y * width + x;
            int index_U = width * height + y / 2 * width / 2 + x / 2;
            int index_V = width * height + width * height / 4 + y / 2 * width / 2 + x / 2;

            // 取出 YUV
            uchar Y = yuv_buffer[index_Y];
            uchar U = yuv_buffer[index_U];
            uchar V = yuv_buffer[index_V];

            // YCbCr420
            int R = Y + 1.402 * (V - 128);
            int G = Y - 0.34413 * (U - 128) - 0.71414*(V - 128);
            int B = Y + 1.772*(U - 128);

            // Y'Cb'Cr'420
//            int R = 1.164 * (Y - 16) + 1.596 * (V - 128);
//            int G = 1.164 * (Y - 16) - 0.813 * (V - 128) - 0.392 * (U - 128);
//            int B = 1.164 * (Y - 16) + 2.017 * (U - 128);

            // 确保取值范围在 0 - 255 中
            R = (R < 0) ? 0 : R;
            G = (G < 0) ? 0 : G;
            B = (B < 0) ? 0 : B;
            R = (R > 255) ? 255 : R;
            G = (G > 255) ? 255 : G;
            B = (B > 255) ? 255 : B;

            rgb_buffer[(y*width + x)*channels + 0] = (uchar)R;
            rgb_buffer[(y*width + x)*channels + 1] = (uchar)G;
            rgb_buffer[(y*width + x)*channels + 2] = (uchar)B;
        }
    }
}



int save_bmp(AVFrame *frame, char * filename)
{
	if (!frame || !filename ){
		return 0;
	}
 
	int frame_channels = 3*sizeof(uchar);
	//写位图数据进文件
	int image_size  = frame->height * frame->width * sizeof(uchar)*frame_channels;
	uchar* rgb_data = (uchar*)malloc(image_size);
    AVFrame2Img(frame, rgb_data);
    
    // printf("__LINE__:%d\n",__LINE__);
    rgbaToBmpFile(filename, rgb_data, frame->width,frame->height);
    // printf("__LINE__  %s:%d\n",__FILE__, __LINE__);
    free(rgb_data);
    return 0;
	// fwrite(rgb_data, image_size, 1, fp);
}